#include <stdio.h>
#include <string.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT *tcti_context = NULL;
TSS2_SYS_CONTEXT *sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;

static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

// Check if the return code indicates success, clean up if not
#define CHECK_PASSED(rc)                                                       \
    {                                                                          \
        if ((rc) != TPM2_RC_SUCCESS) {                                         \
            printf("FAILED! [%s:%u] %s\n", __FUNCTION__, __LINE__,             \
                   Tss2_RC_Decode(rc));                                        \
            cleanup(rc);                                                       \
        } else {                                                               \
            printf("PASSED! [%s:%u]\n", __FUNCTION__, __LINE__);               \
        }                                                                      \
    }

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context,
                         TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size,
                                     tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc =
        Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                               TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush =
                Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n",
                       handle_recorder.handles[i]);
            }
        }
    }
    handle_recorder.count = 0;
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Create a RSA 2048 key
TSS2_RC create_key_rsa_2048(TSS2_SYS_CONTEXT *sys_context, const char *password,
                            TPM2_HANDLE *key_handle, TPM2B_PUBLIC *out_public) {
    // CreatePrimary Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {
        .count = 1, .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPM2B_PUBLIC in_public = {
        .size = 0,
        .publicArea =
            {
                .type = TPM2_ALG_RSA,
                .nameAlg = TPM2_ALG_SHA256,
                .objectAttributes = TPMA_OBJECT_DECRYPT |
                                    TPMA_OBJECT_USERWITHAUTH |
                                    TPMA_OBJECT_SENSITIVEDATAORIGIN,
                .parameters.rsaDetail =
                    {
                        .scheme = {.scheme = TPM2_ALG_OAEP,
                                   .details = {.oaep = {.hashAlg =
                                                            TPM2_ALG_SHA256}}},
                        .keyBits = 2048,
                        .exponent = 0,
                        .symmetric = {.algorithm = TPM2_ALG_NULL},
                    },
            },
    };
    UINT16 pwd_size = strlen(password);
    TPM2B_SENSITIVE_CREATE in_sensitive = {
        .sensitive.userAuth =
            {
                .size = pwd_size,
                .buffer = {},
            },
    };
    memcpy(in_sensitive.sensitive.userAuth.buffer, password, pwd_size);
    TPM2B_DATA outside_info = {
        .size = 0,
        .buffer = {},
    };
    TPML_PCR_SELECTION creation_pcr = {
        .count = 0,
    };

    // CreatePrimary Output Parameters
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc = Tss2_Sys_CreatePrimary(
        sys_context, TPM2_RH_OWNER, &cmd_auths_array, &in_sensitive, &in_public,
        &outside_info, &creation_pcr, key_handle, out_public, &creation_data,
        &creation_hash, &creation_ticket, &name, &rsp_auths_array);
    return rc;
};

TSS2_RC load_pubkey(TSS2_SYS_CONTEXT *sys_context, TPM2B_PUBLIC *in_public,
                    TPM2_HANDLE *key_handle) {

    TSS2_RC rc = Tss2_Sys_LoadExternal(sys_context, NULL, NULL, in_public,
                                       TPM2_RH_NULL, key_handle, NULL, NULL);
    return rc;
}

// RSA Encrypt data using public key
TSS2_RC rsa_pubkey_encrypt(TSS2_SYS_CONTEXT *sys_context, const char *password,
                           TPM2_HANDLE key_handle,
                           TPM2B_PUBLIC_KEY_RSA *in_data,
                           TPM2B_PUBLIC_KEY_RSA *out_data) {
    // Encrypt Input Parameters
    TPMT_RSA_DECRYPT in_scheme = {
        .scheme = TPM2_ALG_OAEP,
        .details = {.oaep = {.hashAlg = TPM2_ALG_SHA256}}};
    TPM2B_DATA label = {0};
    TSS2_RC rc = Tss2_Sys_RSA_Encrypt(sys_context, key_handle, NULL, in_data,
                                      &in_scheme, &label, out_data, NULL);
    return rc;
}

// RSA Decrypt data using private key
TSS2_RC rsa_prvkey_decrypt(TSS2_SYS_CONTEXT *sys_context, const char *password,
                           TPM2_HANDLE key_handle,
                           TPM2B_PUBLIC_KEY_RSA *in_data,
                           TPM2B_PUBLIC_KEY_RSA *out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(password)},
        }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));
    TPMT_RSA_DECRYPT in_scheme = {
        .scheme = TPM2_ALG_OAEP,
        .details = {.oaep = {.hashAlg = TPM2_ALG_SHA256}}};
    TPM2B_DATA label = {0};

    // EncryptDecrypt Output Parameters
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_RSA_Decrypt(sys_context, key_handle, &cmd_auths_array, in_data,
                             &in_scheme, &label, out_data, &rsp_auths_array);
    return rc;
}

int main(int argc, char *argv[]) {

    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_context);
    CHECK_PASSED(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_PASSED(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_PASSED(rc);

    // R: Create a RSA 2048 key
    const char *password = "HelloWorld";
    TPM2_HANDLE key_handle_R;
    TPM2B_PUBLIC out_public_R;
    rc = create_key_rsa_2048(sys_context, password, &key_handle_R,
                             &out_public_R);
    CHECK_PASSED(rc)
    handle_recorder.handles[handle_recorder.count++] =
        key_handle_R; // Record the key handle

    // S: Load public key
    TPM2_HANDLE key_handle_S;
    rc = load_pubkey(sys_context, &out_public_R, &key_handle_S);
    CHECK_PASSED(rc);
    handle_recorder.handles[handle_recorder.count++] =
        key_handle_S; // Record the key handle

    // S: RSA Pubkey Encrypt data
    const char *message = "This is a test message.";
    printf("Plain Text: %s\n", message);
    TPM2B_PUBLIC_KEY_RSA in_encrypt_data = {
        .size = strlen(message),
        .buffer = {0},
    };
    memcpy(in_encrypt_data.buffer, message, in_encrypt_data.size);
    TPM2B_PUBLIC_KEY_RSA out_encrypt_data = {0};
    rc = rsa_pubkey_encrypt(sys_context, NULL, key_handle_S, &in_encrypt_data,
                            &out_encrypt_data);
    CHECK_PASSED(rc);

    // R: RSA Prvkey Decrypt data
    TPM2B_PUBLIC_KEY_RSA out_decrypt_data = {
        sizeof(TPM2B_PUBLIC_KEY_RSA) - 2,
    };
    rc = rsa_prvkey_decrypt(sys_context, password, key_handle_R, &out_encrypt_data,
                            &out_decrypt_data);
    CHECK_PASSED(rc);
    printf("Decrypted Text: %.*s\n", out_decrypt_data.size, out_decrypt_data.buffer);

    // Cleanup resources
    cleanup(rc);

    return 0;
}
